# @author  Sahir Shahryar <sahirshahryar@uga.edu>, Erika Luz, Jamie Monogan
# @since   Monday, January 28, 2018
# @version 3.0 - Streamlined

# Cleanup and options from the original source code
rm(list = ls())
options(scipen = 8)
set.seed(715130425,kind="Mersenne-Twister")

# Load appropriate libraries
library(geoR)
library(spBayes)

# Load dataset, BEF.dat, and create logged form of dependent var
data(BEF.dat)
BEF.dat <- BEF.dat[BEF.dat$ALLBIO02_KGH > 0, ]
BEF.dat$log.bio <- log(BEF.dat$ALLBIO02_KGH * 0.001)

# Crucial to make slicing technique work: reset row names after some extraneous have been removed from BEF.dat
row.names(BEF.dat) <- seq(length = nrow(BEF.dat))

# Bayes model function
bayes_model_function <- function(iterations = 1, sample = 50, 
        dataset = BEF.dat, dv = BEF.dat$log.bio,
        form = ~ELEV + SLOPE + SUM_02_TC1 + SUM_02_TC2 + SUM_02_TC3,
        x = BEF.dat$XUTM, y = BEF.dat$YUTM, 
        low.tau = 0, high.tau = 8,
        low.phi = 50, high.phi = 650,
        slices = 5, cache = NULL) {

    minSampleSize <- length(attr(terms(form), "variables")) + 4
    if (sample < minSampleSize) {
        stop(paste0("bayes_model_function(): Sample size must be at least ", minSampleSize))
    }

    sampleSize <- sample
    bayes.percent <- sampleSize / nrow(dataset)

    # Caching functionality: if multiple iterations are being run, the `cache` parameter
    # can be set to skip repeat runs of krige.bayes on the entire dataset.
    if (is.null(cache)) {
        was.null <- TRUE
        print("*** Running krige.bayes with the full dataset first")
        print(Sys.time())
        bayes.test <- krige.bayes(coords = cbind(x, y), data = dv, locations = NULL,
                model = model.control(trend.d = trend.spatial(form, dataset), 
                    cov.model = "exponential"), 
                prior = prior.control(beta.prior = "flat", tausq.rel.prior = "uniform",
                    tausq.rel.discrete = seq(from = low.tau, to = high.tau, length = 100),
                    phi.discrete = seq(low.phi, high.phi, length = 100)),
                output = output.control(messages = FALSE))
    } else {
        was.null <- FALSE
        bayes.test <- cache
    }

    # Collate results from krige.bayes on the full dataset.
    bayes.results <- bayes.test$posterior$sample
    columns <- dim(bayes.results)[2]
    bayes.whole.percentile <- t(apply(bayes.results, 2, quantile, c(0.5, 0.05, 0.95)))
    whole.bayes.mean.mat <- matrix(apply(bayes.results, 2, mean), ncol = columns, nrow = 1)
    
    # Print results if caching was not used.
    if (was.null) {
        print(Sys.time())
        
        print("*** Bayes whole data summary, posterior and summary", quote = FALSE)
        print(summary(bayes.results))

        print("*** Whole data analysis: median, 5th, and 95th percentile (bayes.whole.percentile)", quote = FALSE)
        print(bayes.whole.percentile)
    }
    
    # Output matrix for iterations
    all.bayes <- matrix(NA, nrow = iterations, ncol = columns)

    # Running the iterations
    for (i in 1:iterations) {
        sample.indices <- sample(x = nrow(dataset), size = sampleSize)
    
        bayes.rand <- dataset[sample.indices, ]
        x.sub <- x[sample.indices]
        y.sub <- y[sample.indices]
        dv.sub <- dv[sample.indices]

        bayes.iter <- krige.bayes(coords = cbind(x.sub, y.sub), data = dv.sub, 
            locations = NULL,
            model = model.control(trend.d = trend.spatial(form, bayes.rand),  
                cov.model = "exponential"),
            prior = prior.control(beta.prior = "flat", tausq.rel.prior = "uniform", 
                tausq.rel.discrete = seq(from = low.tau, to = high.tau, length = 100),
                phi.discrete = seq(low.phi, high.phi, length = 100)),
            output = output.control(messages = FALSE))

        bayes.sub.results <- as.matrix(bayes.iter$posterior$sample)
        all.bayes[i, ] <- apply(bayes.sub.results,2,mean)
    }

    print(paste("***dataset size =", dim(dataset)[1]), quote = FALSE)
    print(paste("***", iterations, "iterations, n =", sampleSize, "and p =", bayes.percent), quote = FALSE)

    # Calculate 50th, 5th, and 95th percentiles of combined matrix
    bayes.all.percentile <- t(apply(all.bayes, 2, quantile, c(0.5, 0.05, 0.95)))
	print("*** Bootstrap results: median, 5th, and 95th percentile (bayes.whole.percentile)", quote = FALSE)
    print(bayes.all.percentile)

    # Combine means of whole dataset with average.bayes.mat
    print("***Bayes mean matrix (average.bayes.mat)", quote = FALSE)
    average.bayes.mat <- matrix(apply(all.bayes, 2, mean), ncol = columns, nrow = 1)
    average.bayes.mat <- rbind(average.bayes.mat, whole.bayes.mean.mat)
    rownames(average.bayes.mat) <- c("all.bayes Means", "Whole Data Set Means")
    print(average.bayes.mat)

    # Find difference
    print("***Whole data set means vs. all.bayes Means difference (bayes.diff)", quote = FALSE)
    bayes.diff <- diff(average.bayes.mat)
    rownames(bayes.diff) <- ("Diff")
    print(bayes.diff)

    # Find PMAD
    print("***Percent mean absolute differences (bayes.pmad):", quote = FALSE)
    iterationMeans <- average.bayes.mat["all.bayes Means", ]
    wholeMeans <- average.bayes.mat["Whole Data Set Means", ]
    bayes.pmad <- 100 * abs((iterationMeans - wholeMeans) / wholeMeans)
    print(bayes.pmad)

    # Print finish
    print(paste("***bayes_model_function finished at", Sys.time()), quote = FALSE)
    print("-----------------------------------------------------------", quote = FALSE)

    return(list(all.bayes = all.bayes,
                bayes.all.percentile = bayes.all.percentile,
                bayes.results = bayes.results,
                bayes.whole.percentile = bayes.whole.percentile,
                average.bayes.mat = average.bayes.mat,
                bayes.diff = bayes.diff,
                bayes.pmad = bayes.pmad,
                cache = bayes.test))
}

# Loop function for Bayes
bayes_iterations <- function(iterations = NULL, sample = 50, cache = NULL) {
    result <- NULL

    if (sample <= 1) {
        percent <- sample * 100
        print(paste("Starting bayes_iterations (", percent, "%),", iterations, 
                    "iterations"), quote = FALSE)
    } else {
        print(paste("Starting bayes_iterations (n =", sample, "),", iterations, 
                    "iterations"), quote = FALSE)
    }

    print(Sys.time())

    for (i in 1:length(iterations)) {
        iteration <- bayes_model_function(iterations = iterations[i], sample = sample, cache = cache)

		#if cache is empty, fill it with the full model from the first run
        if (is.null(cache)) {
            cache <- iteration$cache
        }

		#if there is no result matrix, fill it in
        if (is.null(result)) {
            result <- matrix(NA, nrow = length(iterations), ncol = 9)
        }

        result[i, ] <- iteration$bayes.pmad
    }

    if (sample <= 1) {
        percent <- sample * 100
        print(paste("Finishing bayes_iterations (", percent, "%),", iterations, 
                    "iterations"), quote = FALSE)
    } else {
        print(paste("Finishing bayes_iterations (n =", sample, "),", iterations, 
                    "iterations"), quote = FALSE)
    }

    return(list(result = result, cache = cache))
}


# Exporter function. Writes CSV and plots
export.bayes <- function(plot = NULL, iterations = NULL, sample = 50, directory = NULL) {
    names <- c("beta0", "beta1", "beta2", "beta3", "beta4", "beta5",
               "sigmasq", "phi", "tausq.rel")

    colnames(plot) <- names
    rownames(plot) <- c(iterations)
    
    if (sample <= 1) {
        percent <- sample * 100

        csvFilename <- paste0("bayes.", percent, "pct.csv")
        yLabel <- paste0("bayes.pmad (", percent, "%)")
    } else {
        csvFilename <- paste0("bayes.", sample, "n.csv")
        yLabel <- paste0("bayes.pmad (n = ", sample, ")")
    }

    if (!is.null(directory)) {
        csvFilename <- paste0(directory, "/", csvFilename)
    }

    write.csv(plot, csvFilename, row.names = TRUE)

    for (i in 1:length(names)) {
        if (sample <= 1) {
            pdfFilename <- paste0("bayes.", percent, "pct.", names[i], ".pdf")
        } else {
            pdfFilename <- paste0("bayes.", sample, "n.", names[i], ".pdf")
        }

        if (!is.null(directory)) {
            pdfFilename <- paste0(directory, "/", pdfFilename)
        }

        pdf(pdfFilename,width=4,height=4)
        par(las = 1)
        plot(iterations, plot[, i], xlab = "Iterations", ylab = yLabel, main = names[i], 
             type = "b")
        dev.off()
    }

}

# Looper function
permutation.bayes <- function(iterationClasses = NULL, sampleClasses = NULL,
                              exportDirectory = NULL) {
    cache <- NULL

    print(paste0("Starting program execution at ", Sys.time()), quote = FALSE)

    for (sample in sampleClasses) {
        programRun <- bayes_iterations(iterations = iterationClasses, sample = sample, cache = cache)

        export.bayes(plot = programRun$result, iterations = iterationClasses, sample = sample,
                     directory = exportDirectory)

        if (is.null(cache)) {
            cache <- programRun$cache
        }
    }

    print(paste0("Program execution completed at ", Sys.time()), quote = FALSE)
}

#RUN IT ALL
iterationSizes <- c(1:20, 30, 40, 50)
sampleSizes <- c(50, 150, 250)
#sampleSizes<-c(250)
permutation.bayes(iterationClasses = iterationSizes, sampleClasses = sampleSizes, exportDirectory = "POLSData")
